// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use std::{env, fs, mem, path::PathBuf, process, sync::Arc};

use anyhow::bail;
use crossbeam::queue::SegQueue;
use once_cell::sync::Lazy;
use parking_lot::Mutex;
use tracing::debug;

use xgpu_common::{
    ipc::message::Request,
    recorder::{RecorderConfig, RequestLogger, RequestPlayer},
};

static REQUEST_LOG_FILE_PATH: Lazy<PathBuf> = Lazy::new(|| {
    env::current_dir()
        .unwrap_or_else(|_| env::temp_dir())
        .join(format!("{}_req_log", process::id()))
});
static REQUEST_LOGGER: Lazy<Mutex<Option<RequestLogger>>> = Lazy::new(|| Mutex::new(None));
static REQUEST_PLAYER: Lazy<Mutex<Option<RequestPlayer>>> = Lazy::new(|| Mutex::new(None));

static MINIBATCH_REQUESTS: Lazy<Mutex<Arc<SegQueue<Vec<u8>>>>> =
    Lazy::new(|| Mutex::new(Arc::new(SegQueue::new())));

pub fn initialize() -> anyhow::Result<()> {
    fs::OpenOptions::new()
        .create(true)
        .append(true)
        .open(REQUEST_LOG_FILE_PATH.as_path())?;

    let config = RecorderConfig {
        output_path: REQUEST_LOG_FILE_PATH.as_path().to_owned(),
        buffer_size: 1024,
        max_file_size: 1024 * 1024,
        ..Default::default()
    };

    let mut logger = REQUEST_LOGGER.lock();
    if logger.is_none() {
        *logger = Some(RequestLogger::new(config)?);
        debug!(
            "[Journal] Initializing RequestLogger for log file: {}",
            REQUEST_LOG_FILE_PATH.display()
        );
        return Ok(());
    }

    Ok(())
}

pub fn finalize() {
    let mut logger = REQUEST_LOGGER.lock();

    if let Some(logger) = logger.take() {
        let _ = logger.close();
    }
    *logger = None;

    let _ = fs::remove_file(REQUEST_LOG_FILE_PATH.as_path());
}

pub fn dump_request(req_bytes: &[u8]) -> anyhow::Result<()> {
    let mut logger = REQUEST_LOGGER.lock();
    if let Some(l) = logger.as_mut() {
        l.log_request_serialized(req_bytes.to_vec())?;
        debug!("[Journal] Dumped minibatch request: id={}", req_bytes[0]);
    }

    Ok(())
}

pub fn add_minibatch_request(req_bytes: &[u8]) -> anyhow::Result<()> {
    let queue = {
        let guard = MINIBATCH_REQUESTS.lock();
        guard.clone()
    };
    queue.push(req_bytes.to_owned());

    Ok(())
}

pub fn fetch_minibatch_request() -> anyhow::Result<Option<Vec<u8>>> {
    let queue = {
        let guard = MINIBATCH_REQUESTS.lock();
        guard.clone()
    };

    Ok(queue.pop())
}

pub fn reset_minibatch_requests() -> anyhow::Result<Arc<SegQueue<Vec<u8>>>> {
    let mut guard = MINIBATCH_REQUESTS.lock();
    let old = mem::replace(&mut *guard, Arc::new(SegQueue::new()));

    Ok(old)
}

pub fn request_player_init() -> anyhow::Result<()> {
    let mut player = REQUEST_PLAYER.lock();
    *player = Some(RequestPlayer::from_file(REQUEST_LOG_FILE_PATH.as_path())?);
    debug!(
        "[Journal] Initializing RequestPlayer for log file: {}",
        REQUEST_LOG_FILE_PATH.display()
    );

    Ok(())
}

pub fn load_next_request() -> anyhow::Result<Option<Request<'static>>> {
    let mut player = REQUEST_PLAYER.lock();
    match player.as_mut() {
        Some(p) => match p.next_bytewise()? {
            Some(mut block) => {
                let req: Request = block.bytewise()?;
                Ok(Some(req))
            }
            None => Ok(None),
        },
        None => bail!("RequestPlayer is not initialized"),
    }
}
